import numpy as np
import torch
from scipy.stats import pearsonr

def calculate_initial_values(train_data, aes_model):
    """
    计算每个维度的初始美学感知值。

    参数：
        train_data: DataLoader，训练数据加载器。
        aes_model: nn.Module，美学感知模型。

    返回：
        initial_values: ndarray，包含每个维度初始美学感知值的数组。
    """
    # 用于存储每个维度的能力值累积和计数
    dim_sum = None
    dim_count = None

    # 遍历训练数据
    for img, _, _ in train_data:
        # 使用模型推理美学能力
        with torch.no_grad():
            batch_ability, _ = aes_model(img.cuda())  # 输出形状为 (batch_size, num_dimensions)

        # Min-Max 归一化到 [0.2, 0.8]
        min_val, max_val = batch_ability.min(), batch_ability.max()
        batch_ability = (batch_ability - min_val) * (0.8 - 0.2) / (max_val - min_val) + 0.2

        # 转换到 CPU 并转为 NumPy
        batch_ability = batch_ability.cpu().numpy()

        # 初始化累积和与计数
        if dim_sum is None:
            dim_sum = np.zeros(batch_ability.shape[1])  # num_dimensions
            dim_count = np.zeros(batch_ability.shape[1])

        # 累加每个维度的能力值
        dim_sum += batch_ability.sum(axis=0)  # 按列累加每个维度的分数
        dim_count += batch_ability.shape[0]  # 累加每个维度的样本计数

    # 计算平均值
    initial_values = dim_sum / dim_count
    return initial_values

def evaluate(predictions, targets):
    """
    计算评估指标：MAE, MSE, RMSE。
    Args:
        predictions (torch.Tensor): 预测值。
        targets (torch.Tensor): 真实值。
    Returns:
        tuple: MAE, MSE, RMSE
    """
    mae = torch.mean(torch.abs(predictions - targets)).item()
    mse = torch.mean((predictions - targets) ** 2).item()
    rmse = torch.sqrt(torch.mean((predictions - targets) ** 2)).item()
    return mae, mse, rmse

# 计算美学能力评分提升的奖励
def calculate_accuracy_improvement_reward(previous_mse, current_mse):
    mse_diff = max(current_mse - previous_mse, 0)  # 计算MSE的差值，若当前MSE小于前一MSE，则差值为0
    reward = 1 / (1 + mse_diff)  # 根据MSE的改进计算奖励（MSE越小，奖励越高）
    return reward



# 计算选中维度的相关性奖励
def calculate_relevance_reward(selected_dimensions):
    if selected_dimensions is None or len(selected_dimensions) < 2:
        return 0  # 返回一个合适的默认值

    correlation_sum = 0
    num_pairs = 0
    # 计算所有维度对之间的皮尔逊相关系数
    for i in range(len(selected_dimensions)):
        for j in range(i + 1, len(selected_dimensions)):
            corr, _ = pearsonr(selected_dimensions[i], selected_dimensions[j])  # 计算相关系数
            correlation_sum += corr
            num_pairs += 1

    avg_correlation = correlation_sum / num_pairs if num_pairs > 0 else 0  # 计算所有维度对的平均相关性
    reward = 1 / (1 + (1 - avg_correlation))  # 根据平均相关性计算奖励
    return reward



# 计算维度选择的惩罚
def calculate_penalty(num_selected_dimensions, penalty_factor):
    penalty = num_selected_dimensions * penalty_factor  # 根据选中的维度数量计算惩罚
    return penalty


# 计算总奖励，包含准确性、相关性和惩罚的加权
def calculate_total_reward(previous_mse, current_mse, selected_dimensions, num_selected_dimensions, penalty_factor,
                           w1=1.0, w2=1.0, w3=0.5):
    accuracy_reward = calculate_accuracy_improvement_reward(previous_mse, current_mse)
    relevance_reward = calculate_relevance_reward(selected_dimensions)
    penalty = calculate_penalty(num_selected_dimensions, penalty_factor)

    # 总奖励 = 准确性奖励 + 相关性奖励 - 惩罚
    total_reward = w1 * accuracy_reward + w2 * relevance_reward - w3 * penalty
    return total_reward
